#Copyright 2012 Daniel C.(llaniscudani@gmail.com)

#This file is part of Radius3000.

#Radius3000 is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.

#Radius3000 is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.

#You should have received a copy of the GNU General Public License
#along with Radius3000.  If not, see <http://www.gnu.org/licenses/>.

#mschapv2.py is the MS-CHAP-V2 Library

__version__ = "0.1"

import logging
import hashlib
import random


class MSCHAPv2:
    ####
    # EAP MS-CHAP-V2
    ####
    
    def __init__(self):
        self.authChallenge = b''
        self.packetPeerChallenge = b''
        self.packetNtResponse = b''
        
    
    def challengeHash(self, peerChallenge, authChallenge, userName):
        # source: http://tools.ietf.org/html/rfc2759#section-8.2
        # return only first 8 bytes!!!!!!!!!
        return hashlib.sha1(peerChallenge + authChallenge + userName).digest()[:8]
    
    
    def ntPasswordHash(self, p):
        # input string
        return hashlib.new('md4', p.encode('utf_16le')).digest()
        
        
    def hashNtPasswordHash(self, p):
        # input bytes
        return hashlib.new('md4', p).digest()
    
    
    def generateAuthenticatorResponse(self, password, peerChallenge, ntResponse,
                authChallenge, userName):
        # source: http://tools.ietf.org/html/rfc2759#section-8.7
        
        magic1 = b'\x4d\x61\x67\x69\x63\x20\x73\x65\x72\x76'\
               + b'\x65\x72\x20\x74\x6f\x20\x63\x6c\x69\x65'\
               + b'\x6e\x74\x20\x73\x69\x67\x6e\x69\x6e\x67'\
               + b'\x20\x63\x6f\x6e\x73\x74\x61\x6e\x74'
            
        magic2 = b'\x50\x61\x64\x20\x74\x6f\x20\x6d\x61\x6b'\
               + b'\x65\x20\x69\x74\x20\x64\x6f\x20\x6d\x6f'\
               + b'\x72\x65\x20\x74\x68\x61\x6e\x20\x6f\x6e'\
               + b'\x65\x20\x69\x74\x65\x72\x61\x74\x69\x6f\x6e'
        
        passwordHash = self.ntPasswordHash(password)
        passwordHashHash = self.hashNtPasswordHash(passwordHash)
        
        digest = hashlib.sha1(passwordHashHash + ntResponse + magic1).digest()
        challenge = self.challengeHash(peerChallenge, authChallenge, userName)
        digest = hashlib.sha1(digest + challenge + magic2).hexdigest().upper()
        return ('S=' + digest).encode('utf_8')
    
    
    def checkAuthenticatorResponse(self, password, ntResponse, peerChallenge,
                authChallenge, userName, receivedResponse):
        # source: http://tools.ietf.org/html/rfc2759#section-8.8

        myResponse = self.generateAuthenticatorResponse(password, ntResponse,
                peerChallenge, authChallenge, userName)

        if myResponse == receivedResponse:
            return True
         
         
    def getMasterKey(self, passwordHashHash, ntResponse):
        # http://www.ietf.org/rfc/rfc3079.txt section 3.4
        # return only 16bytes
        magic1 = b'\x54\x68\x69\x73\x20\x69\x73\x20\x74'\
               + b'\x68\x65\x20\x4d\x50\x50\x45\x20\x4d'\
               + b'\x61\x73\x74\x65\x72\x20\x4b\x65\x79' 
        return hashlib.sha1(passwordHashHash + ntResponse + magic1).digest()[:16]


    
    def getAsymmetricStartKey(self, masterKey, sessionKeyLength, isSend, isServer):
        # http://www.ietf.org/rfc/rfc3079.txt section 3.4
        magic2 = b'\x4f\x6e\x20\x74\x68\x65\x20\x63\x6c\x69'\
               + b'\x65\x6e\x74\x20\x73\x69\x64\x65\x2c\x20'\
               + b'\x74\x68\x69\x73\x20\x69\x73\x20\x74\x68'\
               + b'\x65\x20\x73\x65\x6e\x64\x20\x6b\x65\x79'\
               + b'\x3b\x20\x6f\x6e\x20\x74\x68\x65\x20\x73'\
               + b'\x65\x72\x76\x65\x72\x20\x73\x69\x64\x65'\
               + b'\x2c\x20\x69\x74\x20\x69\x73\x20\x74\x68'\
               + b'\x65\x20\x72\x65\x63\x65\x69\x76\x65\x20'\
               + b'\x6b\x65\x79\x2e'
               
        magic3 = b'\x4f\x6e\x20\x74\x68\x65\x20\x63\x6c\x69'\
               + b'\x65\x6e\x74\x20\x73\x69\x64\x65\x2c\x20'\
               + b'\x74\x68\x69\x73\x20\x69\x73\x20\x74\x68'\
               + b'\x65\x20\x72\x65\x63\x65\x69\x76\x65\x20'\
               + b'\x6b\x65\x79\x3b\x20\x6f\x6e\x20\x74\x68'\
               + b'\x65\x20\x73\x65\x72\x76\x65\x72\x20\x73'\
               + b'\x69\x64\x65\x2c\x20\x69\x74\x20\x69\x73'\
               + b'\x20\x74\x68\x65\x20\x73\x65\x6e\x64\x20'\
               + b'\x6b\x65\x79\x2e'
        
        shsPad1 = b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'\
                + b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'\
                + b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'\
                + b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'

        shsPad2 = b'\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2'\
                + b'\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2'\
                + b'\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2'\
                + b'\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2'
        
        if isSend:
            if isServer:
                s = magic3
            else:
                s = magic2
        
        else:
            if isServer:
                s = magic2
            else:
                s = magic3

        return hashlib.sha1(masterKey + shsPad1 + s + shsPad2).digest()[:sessionKeyLength]
    

    def getNewKeyFromSHA(self, startKey, sessionKey):
        sessionKeyLen = len(sessionKey)
        # http://www.ietf.org/rfc/rfc3078.txt section 7.3
        shaPad1 = b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'\
                + b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'\
                + b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'\
                + b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'


        shaPad2 = b'\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2'\
                + b'\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2'\
                + b'\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2'\
                + b'\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2'
        
        return hashlib.sha1(startKey + shaPad1 + sessionKey + shaPad2).digest()[:sessionKeyLen]
      
    
    def sessionKeysGenerator(self, password):
        logging.debug("[MS-CHAP-V2] Generating Session Keys")
        # generator for sendkey rcvkey
        passwordHash = self.ntPasswordHash(password)
        passwordHashHash = self.hashNtPasswordHash(passwordHash)

        masterKey = self.getMasterKey(passwordHashHash, self.packetNtResponse)
        masterSendKey = self.getAsymmetricStartKey(masterKey, 16, True, True)
        masterReceiveKey = self.getAsymmetricStartKey(masterKey, 16, False, True)
        
        sendSessionKey = self.getNewKeyFromSHA(masterSendKey, masterSendKey)
        receiveSessionKey = self.getNewKeyFromSHA(masterReceiveKey, masterReceiveKey)

        return (sendSessionKey, receiveSessionKey)
    
    
    def challenge(self):
        ####
        # mschapv2 challenge packet
        ####
        
        self.authChallenge = random.getrandbits(128).to_bytes(16, 'big')
        # type 26(0x1a)
        p = b'\x1a'
        # opcode 1(challenge) (1byte)
        p += b'\x01'
        # id (1byte) ???RANDOM???
        p += b'\x04'
        # len(2bytes)
        p += (27).to_bytes(2, 'big')
        # value size (0x10 if challenge len is 16)
        p += b'\x10'
        p += self.authChallenge
        # name(min 1 byte) is the "id" of "system"...
        p += b'qwerty'
        
        logging.debug("Sending MS-CHAP-V2 Challenge")
        return p
        
        
    def sucess_request(self, pName):
        ####
        # mschapv2 sucess request packet (srv->cli).
        # client returns (if all ok) a 
        # sucess response packet (srv<-cli)
        ####
        ####
        # packet response (pResponse)
        # 16 octets: Peer-Challenge
        # 8 octets: Reserved, must be zero
        # 24 octets: NT-Response
        # 1 octet : Flags
        ####
    
        # mschapv2 packet
        # type 26(0x1a)
        p = b'\x1a'
        # opcode 3(sucess) (1byte)
        p += b'\x03'
        # id (1byte)
        p += b'\x02'
        # msg len(2bytes)
        # first count byte is opcode
        p += (54).to_bytes(2, 'big')
        # message
        p += self.generateAuthenticatorResponse('aaaa', self.packetPeerChallenge, self.packetNtResponse,
                self.authChallenge, pName)
        # add a tail with a message(and update the length) M=?????
        p += 'M=querty'.encode('utf_8')
        return p


    def parse(self, p):
        logging.debug("[MS-CHAP-V2] Received a MS-CHAP-V2 packet")
        
        #code of type of packet
        p_opcode = p[1]

        if p_opcode == 2:
            logging.debug("[MS-CHAP-V2] Request")
            p_mschapv2id = p[2]
            p_mslen = p[3:5]
            # received a response
            p_vsize = p[5]
            p_resp = p[6:(p_vsize+6)]
            #######
            # packet response:
            # 16 octets: Peer-Challenge 
            # 8 octets: Reserved, must be zero
            # 24 octets: NT-Response
            # 1 octet : Flags
            #######
            p_name = p[(p_vsize+6):]
            
            self.packetPeerChallenge = p_resp[0:16]
            self.packetNtResponse = p_resp[24:48]
            
            print(self.checkAuthenticatorResponse('aaaa', self.packetPeerChallenge, self.packetNtResponse,
                self.authChallenge, p_name, p_resp))
            
            logging.debug("Send MS-CHAP-V2 sucess request packet")
            # send packet request & packet name(user) to sucess request
            # get response and send to socket
            return self.sucess_request(p_name)

            
        elif p_opcode == 3:
            # client return a Sucess packet (all ok)
            #
            if p == b'\x1a\x03':
                logging.debug("[MS-CHAP-V2] Sucess!")
                logging.debug("[MS-CHAP-V2] Auth OK.")
                return True
            
            
        elif p_opcode == 4:
            # client return a Failure packet
            if p == b'\x1a\x04':
                logging.debug("[MS-CHAP-V2] Failure, closing")
            return False



